import os
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.runnables.utils import ConfigurableField
from langchain.prompts import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
)
from langchain.chat_models import init_chat_model

# 设置API密钥环境变量
os.environ["DASHSCOPE_API_KEY"] = "sk-2c036d80d14a443f8769e072bb09a2e1"

# 初始化通义千问模型实例
# 使用 openai 提供商，因为 DashScope 兼容 OpenAI API 格式
qwen_model = init_chat_model(
    "qwen-plus",
    model_provider="openai",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1"
)

# 配置可选模型机制：允许在运行时根据 configurable 字段切换使用的模型
# 默认使用 qwen_model，可通过配置选择其他扩展模型
model = qwen_model.configurable_alternatives(
    ConfigurableField(id="llm"),
    default_key="qwen",
    # 可以在此处添加更多模型选项
)

# 定义聊天提示模板，将用户输入作为查询内容传入
# 使用 HumanMessagePromptTemplate 构建单轮对话格式的消息结构
prompt = ChatPromptTemplate.from_messages(
    [
        HumanMessagePromptTemplate.from_template("{query}"),
    ]
)

# 构建 LangChain 表达式语言 (LCEL) 调用链
# 数据流依次经过：
# 1. 将输入直接传递给 prompt（RunnablePassthrough）
# 2. 填充 prompt 模板中的变量
# 3. 根据配置选择实际执行的模型进行推理
# 4. 使用 StrOutputParser 解析模型输出为字符串结果
chain = (
        {"query": RunnablePassthrough()}
        | prompt
        | model
        | StrOutputParser()
)

try:
    # 执行调用链并指定使用 "qwen" 模型处理请求
    # 输入问题："请自我介绍"
    ret = chain.with_config(configurable={"llm": "qwen"}).invoke("请自我介绍")

    # 输出最终响应结果到控制台
    print(ret)
except Exception as e:
    print(f"调用模型时发生错误: {e}")
    print("请检查网络连接和API密钥配置")